# import bm3d
import numpy
import numpy as np


def bm3d_denoise(image):
    if image.ndim == 3:
        result = np.zeros_like(image)
        for i in range(result.shape[0]):
            result[i] = bm3d.bm3d(image[i], sigma_psd=0.5, stage_arg=bm3d.BM3DStages.ALL_STAGES)
        return result
    elif image.ndim == 2:
        result = bm3d.bm3d(image, sigma_psd=0.5, stage_arg=bm3d.BM3DStages.ALL_STAGES)
        return result
    else:
        raise Exception("Wrong image shape")


def preprocess(img_list):
    if isinstance(img_list, list):
        for i in range(len(img_list)):
            if np.max(img_list[i]) - np.min(img_list[i]) != 0:
                img_list[i] = (img_list[i] - np.min(img_list[i])) / (np.max(img_list[i]) - np.min(img_list[i])) * 255
            else:
                img_list[i] = numpy.zeros_like(img_list[i])
    else:  # numpy array object
        img_list = (img_list - np.min(img_list)) / (np.max(img_list) - np.min(img_list)) * 255
        return img_list
